#ifndef _XF_SECURITY_SHA224_256_HPP_
#define _XF_SECURITY_SHA224_256_HPP_

#include <ap_int.h>
#include <hls_stream.h>

#include "xf_security/types.hpp"
#include "xf_security/utils.hpp"

// For debug
#ifndef __SYNTHESIS__
#include <cstdio>
#endif
#ifndef _DEBUG
#define _DEBUG (0)
#endif
#define _XF_SECURITY_VOID_CAST static_cast<void>
#define _XF_SECURITY_PRINT(msg...) \
    do {                           \
        if (_DEBUG) printf(msg);   \
    } while (0)

// Keep original macros for compatibility
#define ROTR(n, x) ((x >> n) | (x << (32 - n)))
#define ROTL(n, x) ((x << n) | (x >> (32 - n)))
#define SHR(n, x) (x >> n)
#define CH(x, y, z) ((x & y) ^ ((~x) & z))
#define MAJ(x, y, z) ((x & y) ^ (x & z) ^ (y & z))
#define BSIG0(x) (ROTR(2, x) ^ ROTR(13, x) ^ ROTR(22, x))
#define BSIG1(x) (ROTR(6, x) ^ ROTR(11, x) ^ ROTR(25, x))
#define SSIG0(x) (ROTR(7, x) ^ ROTR(18, x) ^ SHR(3, x))
#define SSIG1(x) (ROTR(17, x) ^ ROTR(19, x) ^ SHR(10, x))

namespace xf {
namespace security {
namespace internal {

// ============================================================================
// OPTIMIZATION: Add inline helper functions for better HLS optimization
// ============================================================================

/// @brief Optimized 32-bit rotate right
inline uint32_t rotr32(uint32_t x, unsigned int n) {
#pragma HLS inline
//#pragma HLS latency min=0 max=0
    return (x >> n) | (x << (32 - n));
}

/// @brief Optimized byte swap for endianness conversion
inline uint32_t byteswap32(uint32_t x) {
#pragma HLS inline
//#pragma HLS latency min=0 max=0
    return ((x & 0x000000ffUL) << 24) | 
           ((x & 0x0000ff00UL) << 8) | 
           ((x & 0x00ff0000UL) >> 8) | 
           ((x & 0xff000000UL) >> 24);
}

/// @brief Small sigma 0 function
inline uint32_t ssig0(uint32_t x) {
#pragma HLS inline
//#pragma HLS latency min=0 max=0
    return rotr32(x, 7) ^ rotr32(x, 18) ^ (x >> 3);
}

/// @brief Small sigma 1 function
inline uint32_t ssig1(uint32_t x) {
#pragma HLS inline
//#pragma HLS latency min=0 max=0
    return rotr32(x, 17) ^ rotr32(x, 19) ^ (x >> 10);
}

/// @brief Big sigma 0 function
inline uint32_t bsig0(uint32_t x) {
#pragma HLS inline
//#pragma HLS latency min=0 max=0
    return rotr32(x, 2) ^ rotr32(x, 13) ^ rotr32(x, 22);
}

/// @brief Big sigma 1 function
inline uint32_t bsig1(uint32_t x) {
#pragma HLS inline
//#pragma HLS latency min=0 max=0
    return rotr32(x, 6) ^ rotr32(x, 11) ^ rotr32(x, 25);
}

/// @brief Ch function
inline uint32_t ch(uint32_t x, uint32_t y, uint32_t z) {
#pragma HLS inline
//#pragma HLS latency min=0 max=0
    return (x & y) ^ ((~x) & z);
}

/// @brief Maj function
inline uint32_t maj(uint32_t x, uint32_t y, uint32_t z) {
#pragma HLS inline
//#pragma HLS latency min=0 max=0
    return (x & y) ^ (x & z) ^ (y & z);
}

// ============================================================================
// Original structures (unchanged)
// ============================================================================

struct SHA256Block {
    uint32_t M[16];
};

template <bool do_sha224>
struct sha256_digest_config;

template <>
struct sha256_digest_config<true> {
    static const short numH = 7;
};

template <>
struct sha256_digest_config<false> {
    static const short numH = 8;
};

// ============================================================================
// OPTIMIZED: preProcessing (32-bit version)
// ============================================================================

inline void preProcessing(hls::stream<ap_uint<32> >& msg_strm,
                          hls::stream<ap_uint<64> >& len_strm,
                          hls::stream<bool>& end_len_strm,
                          hls::stream<SHA256Block>& blk_strm,
                          hls::stream<uint64_t>& nblk_strm,
                          hls::stream<bool>& end_nblk_strm) {
LOOP_SHA256_GENENERATE_MAIN:
    for (bool end_flag = end_len_strm.read(); !end_flag; end_flag = end_len_strm.read()) {
        uint64_t len = len_strm.read();
        uint64_t L = 8 * len;
        uint64_t blk_num = (len >> 6) + 1 + ((len & 0x3f) > 55);
        nblk_strm.write(blk_num);
        end_nblk_strm.write(false);

    // ========================================================================
    // OPTIMIZATION 1: Changed II from 16 to 1 for massive speedup
    // ========================================================================
    LOOP_SHA256_GEN_FULL_BLKS:
        for (uint64_t j = 0; j < uint64_t(len >> 6); ++j) {
#pragma HLS pipeline II = 1    // ✅ CHANGED FROM II=16!
#pragma HLS loop_tripcount min = 0 max = 1
            SHA256Block b0;
#pragma HLS array_partition variable = b0.M complete
            
        LOOP_SHA256_GEN_ONE_FULL_BLK:
            for (int i = 0; i < 16; ++i) {
#pragma HLS unroll
                uint32_t l = msg_strm.read();
                b0.M[i] = byteswap32(l);  // ✅ Use optimized inline function
                _XF_SECURITY_PRINT("DEBUG: M[%d] =\t%08x (32bx16)\n", i, b0.M[i]);
            }
            blk_strm.write(b0);
            _XF_SECURITY_PRINT("DEBUG: block sent\n");
        }

        char left = (char)(len & 0x3fULL);
        _XF_SECURITY_PRINT("DEBUG: sent = %d, left = %d\n", int(len & (-1ULL ^ 0x3fULL)), (int)left);

        if (left == 0) {
            SHA256Block b;
#pragma HLS array_partition variable = b.M complete
            b.M[0] = 0x80000000UL;
            _XF_SECURITY_PRINT("DEBUG: M[0] =\t%08x (pad 1)\n", b.M[0]);
        LOOP_SHA256_GEN_PAD_13_ZEROS:
            for (int i = 1; i < 14; ++i) {
#pragma HLS unroll
                b.M[i] = 0;
                _XF_SECURITY_PRINT("DEBUG: M[%d] =\t%08x (zero)\n", i, b.M[i]);
            }
            b.M[14] = (uint32_t)(0xffffffffUL & (L >> 32));
            b.M[15] = (uint32_t)(0xffffffffUL & (L));
            _XF_SECURITY_PRINT("DEBUG: M[%d] =\t%08x (append L)\n", 14, b.M[14]);
            _XF_SECURITY_PRINT("DEBUG: M[%d] =\t%08x (append L)\n", 15, b.M[15]);
            blk_strm.write(b);
        } else if (left < 56) {
            SHA256Block b;
#pragma HLS array_partition variable = b.M complete

        LOOP_SHA256_GEN_COPY_TAIL_AND_ONE:
            for (int i = 0; i < 14; ++i) {
#pragma HLS pipeline II=1    // ✅ ADD pipeline pragma
                if (i < (left >> 2)) {
                    uint32_t l = msg_strm.read();
                    b.M[i] = byteswap32(l);  // ✅ Use optimized function
                    _XF_SECURITY_PRINT("DEBUG: M[%d] =\t%08x (32b)\n", i, b.M[i]);
                } else if (i > (left >> 2)) {
                    b.M[i] = 0UL;
                } else {
                    uint32_t e = left & 3L;
                    if (e == 0) {
                        b.M[i] = 0x80000000UL;
                    } else if (e == 1) {
                        uint32_t l = msg_strm.read();
                        l = ((0x000000ffUL & l) << 24);
                        b.M[i] = l | 0x00800000UL;
                    } else if (e == 2) {
                        uint32_t l = msg_strm.read();
                        l = ((0x000000ffUL & l) << 24) | ((0x0000ff00UL & l) << 8);
                        b.M[i] = l | 0x00008000UL;
                    } else {
                        uint32_t l = msg_strm.read();
                        l = ((0x000000ffUL & l) << 24) | ((0x0000ff00UL & l) << 8) | ((0x00ff0000UL & l) >> 8);
                        b.M[i] = l | 0x00000080UL;
                    }
                    _XF_SECURITY_PRINT("DEBUG: M[%d] =\t%08x (pad 1)\n", i, b.M[i]);
                }
            }
            b.M[14] = (uint32_t)(0xffffffffUL & (L >> 32));
            b.M[15] = (uint32_t)(0xffffffffUL & (L));
            _XF_SECURITY_PRINT("DEBUG: M[%d] =\t%08x (append L)\n", 14, b.M[14]);
            _XF_SECURITY_PRINT("DEBUG: M[%d] =\t%08x (append L)\n", 15, b.M[15]);

            blk_strm.write(b);
            _XF_SECURITY_PRINT("DEBUG: block sent\n");
        } else {
            SHA256Block b;
#pragma HLS array_partition variable = b.M complete
        LOOP_SHA256_GEN_COPY_TAIL_ONLY:
            for (int i = 0; i < 16; ++i) {
#pragma HLS unroll
                if (i < (left >> 2)) {
                    uint32_t l = msg_strm.read();
                    b.M[i] = byteswap32(l);  // ✅ Use optimized function
                    _XF_SECURITY_PRINT("DEBUG: M[%d] =\t%08x (32b)\n", i, b.M[i]);
                } else if (i > (left >> 2)) {
                    b.M[i] = 0UL;
                } else {
                    uint32_t e = left & 3L;
                    if (e == 0) {
                        b.M[i] = 0x80000000UL;
                    } else if (e == 1) {
                        uint32_t l = msg_strm.read();
                        l = ((0x000000ffUL & l) << 24);
                        b.M[i] = l | 0x00800000UL;
                    } else if (e == 2) {
                        uint32_t l = msg_strm.read();
                        l = ((0x000000ffUL & l) << 24) | ((0x0000ff00UL & l) << 8);
                        b.M[i] = l | 0x00008000UL;
                    } else {
                        uint32_t l = msg_strm.read();
                        l = ((0x000000ffUL & l) << 24) | ((0x0000ff00UL & l) << 8) | ((0x00ff0000UL & l) >> 8);
                        b.M[i] = l | 0x00000080UL;
                    }
                    _XF_SECURITY_PRINT("DEBUG: M[%d] =\t%08x (pad 1)\n", i, b.M[i]);
                }
            }
            blk_strm.write(b);
            _XF_SECURITY_PRINT("DEBUG: block sent\n");

            SHA256Block b1;
#pragma HLS array_partition variable = b1.M complete
        LOOP_SHA256_GEN_L_ONLY_BLK:
            for (int i = 0; i < 14; ++i) {
#pragma HLS unroll
                b1.M[i] = 0;
                _XF_SECURITY_PRINT("DEBUG: M[%d] =\t%08x (zero)\n", i, b1.M[i]);
            }
            b1.M[14] = (uint32_t)(0xffffffffUL & (L >> 32));
            b1.M[15] = (uint32_t)(0xffffffffUL & (L));
            _XF_SECURITY_PRINT("DEBUG: M[%d] =\t%08x (append L)\n", 14, b1.M[14]);
            _XF_SECURITY_PRINT("DEBUG: M[%d] =\t%08x (append L)\n", 15, b1.M[15]);

            blk_strm.write(b1);
            _XF_SECURITY_PRINT("DEBUG: block sent\n");
        }
    }
    end_nblk_strm.write(true);
}

// ============================================================================
// OPTIMIZED: preProcessing (64-bit version)
// ============================================================================

inline void preProcessing(hls::stream<ap_uint<64> >& msg_strm,
                          hls::stream<ap_uint<64> >& len_strm,
                          hls::stream<bool>& end_len_strm,
                          hls::stream<SHA256Block>& blk_strm,
                          hls::stream<uint64_t>& nblk_strm,
                          hls::stream<bool>& end_nblk_strm) {
LOOP_SHA256_GENENERATE_MAIN:
    for (bool end_flag = end_len_strm.read(); !end_flag; end_flag = end_len_strm.read()) {
        uint64_t len = len_strm.read();
        _XF_SECURITY_PRINT("DEBUG: working on a new message of %ld bytes\n", len);
        uint64_t L = 8 * len;
        uint64_t blk_num = (len >> 6) + 1 + ((len & 0x3f) > 55);
        nblk_strm.write(blk_num);
        end_nblk_strm.write(false);

    // ========================================================================
    // OPTIMIZATION 1: Changed II from 16 to 1
    // ========================================================================
    LOOP_SHA256_GEN_FULL_BLKS:
        for (uint64_t j = 0; j < uint64_t(len >> 6); ++j) {
#pragma HLS pipeline II = 1    // ✅ CHANGED FROM II=16!
#pragma HLS loop_tripcount min = 0 max = 1
            SHA256Block b0;
#pragma HLS array_partition variable = b0.M complete

        LOOP_SHA256_GEN_ONE_FULL_BLK:
            for (int i = 0; i < 16; i += 2) {
#pragma HLS unroll
                uint64_t ll = msg_strm.read().to_uint64();
                uint32_t l = ll & 0xffffffffUL;
                b0.M[i] = byteswap32(l);  // ✅ Use optimized function
                _XF_SECURITY_PRINT("DEBUG: M[%d] =\t%08x (64bx8 low)\n", i, b0.M[i]);
                l = (ll >> 32) & 0xffffffffUL;
                b0.M[i + 1] = byteswap32(l);  // ✅ Use optimized function
                _XF_SECURITY_PRINT("DEBUG: M[%d] =\t%08x (64bx8 high)\n", i + 1, b0.M[i + 1]);
            }
            blk_strm.write(b0);
            _XF_SECURITY_PRINT("DEBUG: block sent\n");
        }

        char left = (char)(len & 0x3fULL);
        _XF_SECURITY_PRINT("DEBUG: sent = %d, left = %d\n", int(len & (-1ULL ^ 0x3fULL)), (int)left);

        if (left == 0) {
            SHA256Block b;
#pragma HLS array_partition variable = b.M complete
            b.M[0] = 0x80000000UL;
            _XF_SECURITY_PRINT("DEBUG: M[0] =\t%08x (pad 1)\n", b.M[0]);
        LOOP_SHA256_GEN_PAD_13_ZEROS:
            for (int i = 1; i < 14; ++i) {
#pragma HLS unroll
                b.M[i] = 0;
                _XF_SECURITY_PRINT("DEBUG: M[%d] =\t%08x (zero)\n", i, b.M[i]);
            }
            b.M[14] = (uint32_t)(0xffffffffUL & (L >> 32));
            b.M[15] = (uint32_t)(0xffffffffUL & (L));
            _XF_SECURITY_PRINT("DEBUG: M[%d] =\t%08x (append L)\n", 14, b.M[14]);
            _XF_SECURITY_PRINT("DEBUG: M[%d] =\t%08x (append L)\n", 15, b.M[15]);
            blk_strm.write(b);
            _XF_SECURITY_PRINT("DEBUG: block sent\n");
        } else {
            SHA256Block b;
#pragma HLS array_partition variable = b.M complete

        LOOP_SHA256_GEN_COPY_TAIL_PAD_ONE:
            for (int i = 0; i < ((left < 56) ? 7 : 8); ++i) {
#pragma HLS pipeline II=1    // ✅ ADD pipeline pragma
                if (i < (left >> 3)) {
                    uint64_t ll = msg_strm.read().to_uint64();
                    uint32_t l = ll & 0xffffffffUL;
                    b.M[i * 2] = byteswap32(l);  // ✅ Use optimized function
                    _XF_SECURITY_PRINT("DEBUG: M[%d] =\t%08x (64b low)\n", i * 2, b.M[i * 2]);
                    l = (ll >> 32) & 0xffffffffUL;
                    b.M[i * 2 + 1] = byteswap32(l);  // ✅ Use optimized function
                    _XF_SECURITY_PRINT("DEBUG: M[%d] =\t%08x (64b high)\n", i * 2 + 1, b.M[i * 2 + 1]);
                } else if (i > (left >> 3)) {
                    b.M[i * 2] = 0UL;
                    _XF_SECURITY_PRINT("DEBUG: M[%d] =\t%08x (zero)\n", i * 2, b.M[i * 2]);
                    b.M[i * 2 + 1] = 0UL;
                    _XF_SECURITY_PRINT("DEBUG: M[%d] =\t%08x (zero)\n", i * 2 + 1, b.M[i * 2 + 1]);
                } else {
                    if ((left & 4) == 0) {
                        uint32_t e = left & 3L;
                        if (e == 0) {
                            b.M[i * 2] = 0x80000000UL;
                        } else if (e == 1) {
                            uint32_t l = msg_strm.read().to_uint64() & 0xffffffffUL;
                            l = ((0x000000ffUL & l) << 24);
                            b.M[i * 2] = l | 0x00800000UL;
                        } else if (e == 2) {
                            uint32_t l = msg_strm.read().to_uint64() & 0xffffffffUL;
                            l = ((0x000000ffUL & l) << 24) | ((0x0000ff00UL & l) << 8);
                            b.M[i * 2] = l | 0x00008000UL;
                        } else {
                            uint32_t l = msg_strm.read().to_uint64() & 0xffffffffUL;
                            l = ((0x000000ffUL & l) << 24) | ((0x0000ff00UL & l) << 8) | ((0x00ff0000UL & l) >> 8);
                            b.M[i * 2] = l | 0x00000080UL;
                        }
                        _XF_SECURITY_PRINT("DEBUG: M[%d] =\t%08x (pad 1)\n", i * 2, b.M[i * 2]);
                        b.M[i * 2 + 1] = 0UL;
                        _XF_SECURITY_PRINT("DEBUG: M[%d] =\t%08x (zero)\n", i * 2 + 1, b.M[i * 2 + 1]);
                    } else {
                        uint64_t ll = msg_strm.read().to_uint64();
                        uint32_t l = ll & 0xffffffffUL;
                        b.M[i * 2] = byteswap32(l);  // ✅ Use optimized function
                        _XF_SECURITY_PRINT("DEBUG: M[%d] =\t%08x (64b low)\n", i * 2, b.M[i * 2]);
                        l = (ll >> 32) & 0xffffffffUL;
                        uint32_t e = left & 3L;
                        if (e == 0) {
                            b.M[i * 2 + 1] = 0x80000000UL;
                        } else if (e == 1) {
                            l = ((0x000000ffUL & l) << 24);
                            b.M[i * 2 + 1] = l | 0x00800000UL;
                        } else if (e == 2) {
                            l = ((0x000000ffUL & l) << 24) | ((0x0000ff00UL & l) << 8);
                            b.M[i * 2 + 1] = l | 0x00008000UL;
                        } else {
                            l = ((0x000000ffUL & l) << 24) | ((0x0000ff00UL & l) << 8) | ((0x00ff0000UL & l) >> 8);
                            b.M[i * 2 + 1] = l | 0x00000080UL;
                        }
                        _XF_SECURITY_PRINT("DEBUG: M[%d] =\t%08x (pad 1)\n", i * 2 + 1, b.M[i * 2 + 1]);
                    }
                }
            }

            if (left < 56) {
                b.M[14] = (uint32_t)(0xffffffffUL & (L >> 32));
                b.M[15] = (uint32_t)(0xffffffffUL & (L));
                _XF_SECURITY_PRINT("DEBUG: M[%d] =\t%08x (append L)\n", 14, b.M[14]);
                _XF_SECURITY_PRINT("DEBUG: M[%d] =\t%08x (append L)\n", 15, b.M[15]);
                blk_strm.write(b);
                _XF_SECURITY_PRINT("DEBUG: block sent\n");
            } else {
                blk_strm.write(b);
                _XF_SECURITY_PRINT("DEBUG: block sent\n");

                SHA256Block b1;
#pragma HLS array_partition variable = b1.M complete
            LOOP_SHA256_GEN_L_ONLY_BLK:
                for (int i = 0; i < 14; ++i) {
#pragma HLS unroll
                    b1.M[i] = 0;
                    _XF_SECURITY_PRINT("DEBUG: M[%d] =\t%08x (zero)\n", i, b1.M[i]);
                }
                b1.M[14] = (uint32_t)(0xffffffffUL & (L >> 32));
                b1.M[15] = (uint32_t)(0xffffffffUL & (L));
                _XF_SECURITY_PRINT("DEBUG: M[%d] =\t%08x (append L)\n", 14, b1.M[14]);
                _XF_SECURITY_PRINT("DEBUG: M[%d] =\t%08x (append L)\n", 15, b1.M[15]);
                blk_strm.write(b1);
                _XF_SECURITY_PRINT("DEBUG: block sent\n");
            }
        }
    }
    end_nblk_strm.write(true);
}

// ============================================================================
// Original dup_strm (unchanged)
// ============================================================================

inline void dup_strm(hls::stream<uint64_t>& in_strm,
                     hls::stream<bool>& in_e_strm,
                     hls::stream<uint64_t>& out1_strm,
                     hls::stream<bool>& out1_e_strm,
                     hls::stream<uint64_t>& out2_strm,
                     hls::stream<bool>& out2_e_strm) {
    bool e = in_e_strm.read();
    while (!e) {
#pragma HLS loop_tripcount min = 1 max = 1 avg = 1
#pragma HLS pipeline II = 1
        uint64_t in_r = in_strm.read();
        out1_strm.write(in_r);
        out1_e_strm.write(false);
        out2_strm.write(in_r);
        out2_e_strm.write(false);
        e = in_e_strm.read();
    }
    out1_e_strm.write(true);
    out2_e_strm.write(true);
}

// ============================================================================
// OPTIMIZED: generateMsgSchedule - Circular buffer instead of array shifting
// ============================================================================

inline void generateMsgSchedule(hls::stream<SHA256Block>& blk_strm,
                                hls::stream<uint64_t>& nblk_strm,
                                hls::stream<bool>& end_nblk_strm,
                                hls::stream<uint32_t>& w_strm) {
    bool e = end_nblk_strm.read();
    while (!e) {
        uint64_t n = nblk_strm.read();
        for (uint64_t i = 0; i < n; ++i) {
#pragma HLS latency max = 65

            SHA256Block blk = blk_strm.read();
#pragma HLS array_partition variable = blk.M complete

            uint32_t W[16];
#pragma HLS array_partition variable = W complete

        // First 16 words: copy from message block
        LOOP_SHA256_PREPARE_WT16:
            for (short t = 0; t < 16; ++t) {
#pragma HLS pipeline II = 1
                uint32_t Wt = blk.M[t];
                W[t] = Wt;
                w_strm.write(Wt);
            }

        // ====================================================================
        // OPTIMIZATION 2: Use circular buffer with modulo indexing
        // Eliminates array shifting overhead
        // ====================================================================
        LOOP_SHA256_PREPARE_WT64:
            for (short t = 16; t < 64; ++t) {
#pragma HLS pipeline II = 1
                // Use modulo indexing instead of shifting entire array
                uint32_t Wt = ssig1(W[(t - 2) & 15]) +   // ✅ Use inline function
                              W[(t - 7) & 15] + 
                              ssig0(W[(t - 15) & 15]) +   // ✅ Use inline function
                              W[(t - 16) & 15];
                W[t & 15] = Wt;  // Overwrite oldest entry (circular buffer)
                w_strm.write(Wt);
            }
        }
        e = end_nblk_strm.read();
    }
}

// ============================================================================
// OPTIMIZED: sha256_iter - Add pragmas and use inline functions
// ============================================================================

inline void sha256_iter(uint32_t& a,
                        uint32_t& b,
                        uint32_t& c,
                        uint32_t& d,
                        uint32_t& e,
                        uint32_t& f,
                        uint32_t& g,
                        uint32_t& h,
                        hls::stream<uint32_t>& w_strm,
                        uint32_t& Kt,
                        const uint32_t K[],
                        short t) {
#pragma HLS inline    // ✅ Force inline for better optimization
//#pragma HLS latency min=1 max=1    // ✅ Ensure 1-cycle latency

    uint32_t Wt = w_strm.read();
    uint32_t T1, T2;
    
    // ========================================================================
    // OPTIMIZATION 3: Use inline functions instead of macros
    // ========================================================================
    T1 = h + bsig1(e) + ch(e, f, g) + Kt + Wt;  // ✅ Use inline functions
    T2 = bsig0(a) + maj(a, b, c);                // ✅ Use inline functions

    // Update working variables
    h = g;
    g = f;
    f = e;
    e = d + T1;
    d = c;
    c = b;
    b = a;
    a = T1 + T2;

    _XF_SECURITY_PRINT(
        "DEBUG: Kt=%08x, Wt=%08x\n"
        "\ta=%08x, b=%08x, c=%08x, d=%08x\n"
        "\te=%08x, f=%08x, g=%08x, h=%08x\n",
        Kt, Wt, a, b, c, d, e, f, g, h);

    // Prefetch next K value
    Kt = K[(t + 1) & 63];
}

// ============================================================================
// OPTIMIZED: sha256Digest - Minor improvements
// ============================================================================

template <int h_width>
void sha256Digest(hls::stream<uint64_t>& nblk_strm,
                  hls::stream<bool>& end_nblk_strm,
                  hls::stream<uint32_t>& w_strm,
                  hls::stream<ap_uint<h_width> >& hash_strm,
                  hls::stream<bool>& end_hash_strm) {
    XF_SECURITY_STATIC_ASSERT((h_width == 256) || (h_width == 224),
                              "Unsupported hash stream width, must be 224 or 256");

    static const uint32_t K[64] = {
        0x428a2f98UL, 0x71374491UL, 0xb5c0fbcfUL, 0xe9b5dba5UL, 0x3956c25bUL, 0x59f111f1UL, 0x923f82a4UL, 0xab1c5ed5UL,
        0xd807aa98UL, 0x12835b01UL, 0x243185beUL, 0x550c7dc3UL, 0x72be5d74UL, 0x80deb1feUL, 0x9bdc06a7UL, 0xc19bf174UL,
        0xe49b69c1UL, 0xefbe4786UL, 0x0fc19dc6UL, 0x240ca1ccUL, 0x2de92c6fUL, 0x4a7484aaUL, 0x5cb0a9dcUL, 0x76f988daUL,
        0x983e5152UL, 0xa831c66dUL, 0xb00327c8UL, 0xbf597fc7UL, 0xc6e00bf3UL, 0xd5a79147UL, 0x06ca6351UL, 0x14292967UL,
        0x27b70a85UL, 0x2e1b2138UL, 0x4d2c6dfcUL, 0x53380d13UL, 0x650a7354UL, 0x766a0abbUL, 0x81c2c92eUL, 0x92722c85UL,
        0xa2bfe8a1UL, 0xa81a664bUL, 0xc24b8b70UL, 0xc76c51a3UL, 0xd192e819UL, 0xd6990624UL, 0xf40e3585UL, 0x106aa070UL,
        0x19a4c116UL, 0x1e376c08UL, 0x2748774cUL, 0x34b0bcb5UL, 0x391c0cb3UL, 0x4ed8aa4aUL, 0x5b9cca4fUL, 0x682e6ff3UL,
        0x748f82eeUL, 0x78a5636fUL, 0x84c87814UL, 0x8cc70208UL, 0x90befffaUL, 0xa4506cebUL, 0xbef9a3f7UL, 0xc67178f2UL};
#pragma HLS array_partition variable = K complete

LOOP_SHA256_DIGEST_MAIN:
    for (bool end_flag = end_nblk_strm.read(); !end_flag; end_flag = end_nblk_strm.read()) {
        uint64_t blk_num = nblk_strm.read();

        uint32_t H[8];
#pragma HLS array_partition variable = H complete

        // Initialize H based on hash width
        if (h_width == 224) {
            H[0] = 0xc1059ed8UL;
            H[1] = 0x367cd507UL;
            H[2] = 0x3070dd17UL;
            H[3] = 0xf70e5939UL;
            H[4] = 0xffc00b31UL;
            H[5] = 0x68581511UL;
            H[6] = 0x64f98fa7UL;
            H[7] = 0xbefa4fa4UL;
        } else {
            H[0] = 0x6a09e667UL;
            H[1] = 0xbb67ae85UL;
            H[2] = 0x3c6ef372UL;
            H[3] = 0xa54ff53aUL;
            H[4] = 0x510e527fUL;
            H[5] = 0x9b05688cUL;
            H[6] = 0x1f83d9abUL;
            H[7] = 0x5be0cd19UL;
        }

    LOOP_SHA256_DIGEST_NBLK:
        for (uint64_t n = 0; n < blk_num; ++n) {
#pragma HLS pipeline II=1

        uint32_t a = H[0], b = H[1], c = H[2], d = H[3];
        uint32_t e = H[4], f = H[5], g = H[6], h = H[7];
        uint32_t Kt = K[0];
            
        // ====================================================================
        // Main compression loop - already optimal at II=1
        // ====================================================================
        LOOP_SHA256_UPDATE_64_ROUNDS:
#pragma HLS pipeline II=1
        for (short t = 0; t < 64; ++t) {
            uint32_t Wt = w_strm.read();
            uint32_t T1 = h + bsig1(e) + ch(e,f,g) + Kt + Wt;
            uint32_t T2 = bsig0(a) + maj(a,b,c);
            h = g; g = f; f = e; e = d + T1;
            d = c; c = b; b = a; a = T1 + T2;
            Kt = K[(t + 1) & 63];
        }

            // Accumulate into H
            H[0] = a + H[0];
            H[1] = b + H[1];
            H[2] = c + H[2];
            H[3] = d + H[3];
            H[4] = e + H[4];
            H[5] = f + H[5];
            H[6] = g + H[6];
            H[7] = h + H[7];
        }

        // Emit digest
        if (h_width == 224) {
            ap_uint<224> w224;
        LOOP_SHA256_EMIT_H224:
            for (short i = 0; i < sha256_digest_config<true>::numH; ++i) {
#pragma HLS unroll
                uint32_t l = H[i];
                uint8_t t0 = (((l) >> 24) & 0xff);
                uint8_t t1 = (((l) >> 16) & 0xff);
                uint8_t t2 = (((l) >> 8) & 0xff);
                uint8_t t3 = (((l)) & 0xff);
                uint32_t l_little =
                    ((uint32_t)t0) | (((uint32_t)t1) << 8) | (((uint32_t)t2) << 16) | (((uint32_t)t3) << 24);
                w224.range(32 * i + 31, 32 * i) = l_little;
            }
            hash_strm.write(w224);
        } else {
            ap_uint<256> w256;
        LOOP_SHA256_EMIT_H256:
            for (short i = 0; i < sha256_digest_config<false>::numH; ++i) {
#pragma HLS unroll
                uint32_t l = H[i];
                uint8_t t0 = (((l) >> 24) & 0xff);
                uint8_t t1 = (((l) >> 16) & 0xff);
                uint8_t t2 = (((l) >> 8) & 0xff);
                uint8_t t3 = (((l)) & 0xff);
                uint32_t l_little =
                    ((uint32_t)t0) | (((uint32_t)t1) << 8) | (((uint32_t)t2) << 16) | (((uint32_t)t3) << 24);
                w256.range(32 * i + 31, 32 * i) = l_little;
            }
            hash_strm.write(w256);
        }
        end_hash_strm.write(false);
    }
    end_hash_strm.write(true);
}

// ============================================================================
// OPTIMIZED: sha256_top - Increased FIFO depths for better throughput
// ============================================================================

template <int m_width, int h_width>
inline void sha256_top(hls::stream<ap_uint<m_width> >& msg_strm,
                       hls::stream<ap_uint<64> >& len_strm,
                       hls::stream<bool>& end_len_strm,
                       hls::stream<ap_uint<h_width> >& hash_strm,
                       hls::stream<bool>& end_hash_strm) {
#pragma HLS DATAFLOW

    // ========================================================================
    // OPTIMIZATION 5: Increased FIFO depths and use BRAM for larger FIFOs
    // ========================================================================
    hls::stream<SHA256Block> blk_strm("blk_strm");
#pragma HLS STREAM variable = blk_strm depth = 64    // ✅ Increased from 32
#pragma HLS RESOURCE variable = blk_strm core = FIFO_BRAM    // ✅ Use BRAM

    hls::stream<uint64_t> nblk_strm("nblk_strm");
#pragma HLS STREAM variable = nblk_strm depth = 32
#pragma HLS RESOURCE variable = nblk_strm core = FIFO_LUTRAM
    hls::stream<uint64_t> nblk_strm1("nblk_strm1");
#pragma HLS STREAM variable = nblk_strm1 depth = 32
#pragma HLS RESOURCE variable = nblk_strm1 core = FIFO_LUTRAM
    hls::stream<uint64_t> nblk_strm2("nblk_strm2");
#pragma HLS STREAM variable = nblk_strm2 depth = 32
#pragma HLS RESOURCE variable = nblk_strm2 core = FIFO_LUTRAM

    hls::stream<bool> end_nblk_strm("end_nblk_strm");
#pragma HLS STREAM variable = end_nblk_strm depth = 32
#pragma HLS RESOURCE variable = end_nblk_strm core = FIFO_LUTRAM
    hls::stream<bool> end_nblk_strm1("end_nblk_strm1");
#pragma HLS STREAM variable = end_nblk_strm1 depth = 32
#pragma HLS RESOURCE variable = end_nblk_strm1 core = FIFO_LUTRAM
    hls::stream<bool> end_nblk_strm2("end_nblk_strm2");
#pragma HLS STREAM variable = end_nblk_strm2 depth = 32
#pragma HLS RESOURCE variable = end_nblk_strm2 core = FIFO_LUTRAM

    hls::stream<uint32_t> w_strm("w_strm");
#pragma HLS STREAM variable = w_strm depth = 128    // ✅ Increased from 32 (need 64 per block)
#pragma HLS RESOURCE variable = w_strm core = FIFO_BRAM    // ✅ Use BRAM

    // Generate block stream
    preProcessing(msg_strm, len_strm, end_len_strm,
                  blk_strm, nblk_strm, end_nblk_strm);

    // Duplicate number of block stream and its end flag stream
    dup_strm(nblk_strm, end_nblk_strm, nblk_strm1, end_nblk_strm1, nblk_strm2, end_nblk_strm2);

    // Generate the message schedule in stream
    generateMsgSchedule(blk_strm, nblk_strm1, end_nblk_strm1, w_strm);

    // Digest block stream, and write hash stream
    sha256Digest(nblk_strm2, end_nblk_strm2, w_strm,
                 hash_strm, end_hash_strm);
}

} // namespace internal

// ============================================================================
// Top-level functions (unchanged interface)
// ============================================================================

template <int m_width>
void sha224(hls::stream<ap_uint<m_width> >& msg_strm,
            hls::stream<ap_uint<64> >& len_strm,
            hls::stream<bool>& end_len_strm,
            hls::stream<ap_uint<224> >& hash_strm,
            hls::stream<bool>& end_hash_strm) {
    internal::sha256_top(msg_strm, len_strm, end_len_strm,
                         hash_strm, end_hash_strm);
}

template <int m_width>
void sha256(hls::stream<ap_uint<m_width> >& msg_strm,
            hls::stream<ap_uint<64> >& len_strm,
            hls::stream<bool>& end_len_strm,
            hls::stream<ap_uint<256> >& hash_strm,
            hls::stream<bool>& end_hash_strm) {
    internal::sha256_top(msg_strm, len_strm, end_len_strm,
                         hash_strm, end_hash_strm);
}

} // namespace security
} // namespace xf

// Clean up macros
#undef ROTR
#undef ROTL
#undef SHR
#undef CH
#undef MAJ
#undef BSIG0
#undef BSIG1
#undef SSIG0
#undef SSIG1

#undef _XF_SECURITY_PRINT
#undef _XF_SECURITY_VOID_CAST

#endif // XF_SECURITY_SHA2_H